#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include "sm4_1.h"
#define KEY_SIZE 16  // 定义密钥长度为16字节
#define BLOCK_SIZE 16  // SM4 block size is 16 bytes

void pkcs7_padding(unsigned char* input, int input_length, unsigned char* output, int* output_length) {
    int padding_length = BLOCK_SIZE - (input_length % BLOCK_SIZE);
    *output_length = input_length + padding_length;

    // Copy input to output and apply PKCS#7 padding
    memcpy(output, input, input_length);
    memset(output + input_length, padding_length, padding_length);
}

void pkcs7_unpadding(unsigned char* input, int input_length, unsigned char* output, int* output_length) {
    // Get padding value
    int padding_value = input[input_length - 1];
    *output_length = input_length - padding_value;
    memcpy(output, input, *output_length);
}

int main(int argc, char* argv[]) {
    if (argc != 2) {
        fprintf(stderr, "Usage: %s <plaintext>\n", argv[0]);
        return 1;
    }

    // Self-check
    if (SM4_SelfCheck() == 0) {
        printf("自检成功！\n");
    }
    else {
        printf("自检失败！\n");
        return 1; // Self-check failed, exit program
    }

    // Generate random key
    unsigned char key[KEY_SIZE];
    srand((unsigned int)time(NULL)); // 使用当前时间作为随机数种子

    for (int i = 0; i < KEY_SIZE; i++) {
        key[i] = rand() % 256; // 生成0-255之间的随机数
    }

    // Output generated key
    printf("生成的随机密钥: ");
    for (int i = 0; i < KEY_SIZE; i++) {
        printf("%02x", key[i]);  // Hexadecimal output without space
    }
    printf("\n");

    // Prepare plaintext
    const char* plaintext_input = argv[1];
    unsigned char plaintext[BLOCK_SIZE] = { 0 }; // Buffer for plaintext
    unsigned char padded_plaintext[BLOCK_SIZE] = { 0 }; // Buffer for padded plaintext
    unsigned char ciphertext[BLOCK_SIZE] = { 0 }; // Buffer for ciphertext
    unsigned char decrypted[BLOCK_SIZE] = { 0 }; // Buffer for decrypted text
    int padded_length = 0;

    // Apply PKCS#7 padding to the plaintext
    pkcs7_padding((unsigned char*)plaintext_input, strlen(plaintext_input), padded_plaintext, &padded_length);

    // Encrypt plaintext
    SM4_Encrypt(key, padded_plaintext, ciphertext);
    printf("明文 (%s) 的密文: ", plaintext_input);
    for (int i = 0; i < BLOCK_SIZE; i++) {
        printf("%02x", ciphertext[i]);  // Hexadecimal output without space
    }
    printf("\n");

    // Decrypt ciphertext
    SM4_Decrypt(key, ciphertext, decrypted);

    // Remove PKCS#7 padding from decrypted text
    int decrypted_length = 0;
    pkcs7_unpadding(decrypted, BLOCK_SIZE, plaintext, &decrypted_length);
    plaintext[decrypted_length] = '\0'; // Null-terminate the decrypted plaintext for printing

    printf("解密后的文本: %s\n", plaintext);

    return 0;
}
